package us.luosl.hadoop.hive.udaf.generic;

import org.apache.hadoop.hive.ql.metadata.HiveException;
import org.apache.hadoop.hive.ql.parse.SemanticException;
import org.apache.hadoop.hive.ql.udf.generic.AbstractGenericUDAFResolver;
import org.apache.hadoop.hive.ql.udf.generic.GenericUDAFEvaluator;
import org.apache.hadoop.hive.serde2.objectinspector.*;
import org.apache.hadoop.hive.serde2.objectinspector.primitive.*;
import org.apache.hadoop.hive.serde2.typeinfo.PrimitiveTypeInfo;
import org.apache.hadoop.hive.serde2.typeinfo.TypeInfo;
import org.apache.hadoop.io.DoubleWritable;
import org.apache.hadoop.io.LongWritable;

import java.util.ArrayList;
import java.util.List;

/**
 * 求平均值
  add jar /mr-1.0-SNAPSHOT.jar;
  CREATE TEMPORARY FUNCTION my_avg AS 'us.luosl.hadoop.hive.udaf.generic.GenericUDAFAvg';
  select my_avg(yka055) from kc21k1;

 *
 * Created by luosl on 2018/9/4.
 */
public class GenericUDAFAvg extends AbstractGenericUDAFResolver {

    @Override
    public GenericUDAFEvaluator getEvaluator(TypeInfo[] info) throws SemanticException {
        if(info.length != 1){
            throw new SemanticException("该函数只接受一个参数!");
        }
        if (info[0].getCategory() != ObjectInspector.Category.PRIMITIVE) {
            throw new SemanticException("该函数double类型参数!");
        }
        switch (((PrimitiveTypeInfo) info[0]).getPrimitiveCategory()) {
            case FLOAT:
            case DOUBLE:
                return new AvgUDAFEvaluator();
            default:
                throw new SemanticException("该函数double类型参数!");
        }
    }

    /**
     * 定义聚合 buffer
     */
    public static class AvgBuffer extends GenericUDAFEvaluator.AbstractAggregationBuffer {
        private double sum;
        private long count;

        /**
         * 估算 buffer所占内存
         * @return
         */
        @Override
        public int estimate() {
            return 32;
        }
    }

    public static class AvgUDAFEvaluator extends GenericUDAFEvaluator{
        // PARTIAL1 COMPLETE 阶段的输入
        private DoubleObjectInspector inputOI;

        // PARTIAL2 FINAL 阶段的输入
        private StructObjectInspector soi;
        private StructField countField;
        private StructField sumField;
        private LongObjectInspector countOI;
        private DoubleObjectInspector sumOI;

        private Object[] resultArr;

        // 定义输出

        @Override
        public ObjectInspector init(Mode m, ObjectInspector[] parameters) throws HiveException {
            super.init(m, parameters);
            // 定义输入
            if(m == Mode.PARTIAL1 || m == Mode.COMPLETE){
                inputOI = (DoubleObjectInspector) parameters[0];
            }else{
                soi = (StructObjectInspector) parameters[0];
                countField = soi.getStructFieldRef("count");
                sumField = soi.getStructFieldRef("sum");
                countOI = (LongObjectInspector) countField.getFieldObjectInspector();
                sumOI = (DoubleObjectInspector) sumField.getFieldObjectInspector();
            }
            // 定义输出
            if(m == Mode.PARTIAL1 || m == Mode.PARTIAL2){
                List<ObjectInspector> fois = new ArrayList<ObjectInspector>();
                List<String> fnames = new ArrayList<String>();
                fois.add(PrimitiveObjectInspectorFactory.writableLongObjectInspector);
                fois.add(PrimitiveObjectInspectorFactory.writableDoubleObjectInspector);
                fnames.add("count");
                fnames.add("sum");
                return ObjectInspectorFactory.getStandardStructObjectInspector(fnames, fois);
            }else{
                return PrimitiveObjectInspectorFactory.writableDoubleObjectInspector;
            }
        }

        @Override
        public AggregationBuffer getNewAggregationBuffer() throws HiveException {
            return new AvgBuffer();
        }

        @Override
        public void reset(AggregationBuffer agg) throws HiveException {
            AvgBuffer avgAgg = (AvgBuffer) agg;
            avgAgg.count = 0L;
            avgAgg.sum = 0D;
        }

        @Override
        public void iterate(AggregationBuffer agg, Object[] parameters) throws HiveException {
            AvgBuffer avgAgg = (AvgBuffer) agg;
            if(null != parameters[0]){
                double value = inputOI.get(parameters[0]);
                avgAgg.sum += value;
            }
            avgAgg.count += 1;
        }

        @Override
        public Object terminatePartial(AggregationBuffer agg) throws HiveException {
            AvgBuffer avgAgg = (AvgBuffer) agg;
            resultArr = new Object[2];
            resultArr[0] = new LongWritable(avgAgg.count);
            resultArr[1] = new DoubleWritable(avgAgg.sum);
            return resultArr;
        }

        @Override
        public void merge(AggregationBuffer agg, Object partial) throws HiveException {
            AvgBuffer avgAgg = (AvgBuffer) agg;
            Object countData = soi.getStructFieldData(partial, countField);
            Object sumData = soi.getStructFieldData(partial, sumField);
            avgAgg.count += countOI.get(countData);
            avgAgg.sum += sumOI.get(sumData);
        }

        @Override
        public Object terminate(AggregationBuffer agg) throws HiveException {
            AvgBuffer avgAgg = (AvgBuffer) agg;
            if(avgAgg.count == 0) {
                return null;
            }
            return new DoubleWritable(avgAgg.sum / avgAgg.count);
        }
    }
}
